import os
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from FSdata.FSdataset import FSdata, collate_fn, attr2catidx_map, idx2attr_map
import torch
import torch.utils.data as torchdata
from models.resnet import resnet50_cat
from FSdata.FSaug import  *
from utils.predicting import predict,gen_submission
from models.NASnet import NASnet_cat, NASnet_cat_dilate,NASnet_cat_dilate2
from utils.preprocessing import join_path_to_df

class FSAugVal(object):
    def __init__(self):
        self.augment = Compose([
            ExpandBorder(select=[0,5, 6, 7], mode='constant',resize=True,size=(336,336)),
            UpperCrop(size=(368, 368), select=[1,2, 3, 4]),
            Resize(size=(368, 368), select=[1, 2, 3, 4]),
            Resize(size=(336, 336), select=[0, 5, 6, 7]),
            Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ])
    def __call__(self, image,attr_idx):
        return self.augment(image,attr_idx)

os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"
rawdata_root = '/media/gserver/data/FashionAI'

round1_df = pd.read_csv(os.path.join(rawdata_root,'round1/base/Annotations/label.csv'),
                        header=None, names=['ImageName', 'AttrKey', 'AttrValues'])
round1_df = join_path_to_df(round1_df, rawdata_root, 'round1/base')


round2_df = pd.read_csv(os.path.join(rawdata_root,'round2/train/Annotations/label.csv'),
                        header=None, names=['ImageName', 'AttrKey', 'AttrValues'])
round2_df = join_path_to_df(round2_df, rawdata_root, 'round2/train')

extra_df = pd.read_csv(os.path.join(rawdata_root,'round2/round2_data_add_skirt_legth.txt'),
                        header=None, names=['ImageName', 'AttrKey', 'AttrValues'])
extra_df = join_path_to_df(extra_df, rawdata_root, 'round1/web')


round2_train_pd, val_pd = train_test_split(round2_df, test_size=0.1, random_state=37,
                                    stratify=round2_df['AttrKey'])


# select part
select_AttrIdx = [1,2,3,4]
select_AttrKey = [idx2attr_map[x] for x in select_AttrIdx]
val_pd = val_pd[val_pd['AttrKey'].apply(lambda x: True if x in select_AttrKey else False)]


print val_pd.shape
print val_pd['AttrKey'].value_counts()

data_set = {}
data_set['val'] = FSdata(anno_pd=val_pd,
                         transforms=FSAugVal(),
                         select=select_AttrIdx
                         )

data_loader = {}
data_loader['val'] = torchdata.DataLoader(data_set['val'], batch_size=8, num_workers=4,
                                          shuffle=False, pin_memory=True,collate_fn=collate_fn)


# model prepare
model_name = 'NASnet_r2_9625(dpred2)'
resume = '/media/gserver/extra/FashionAI/round2/NAS[1234][r2]/SWA[12-13].pth'

model = NASnet_cat_dilate2(pretrained=True, num_classes=data_set['val'].num_classes)

model = torch.nn.DataParallel(model)
print('resuming finetune from %s'%resume)
model.load_state_dict(torch.load(resume))
model = model.cuda()

# predict
if not os.path.exists('./val_pred2/csv'):
    os.makedirs('./val_pred2/csv')
if not os.path.exists('./val_pred2/part_csv'):
    os.makedirs('./val_pred2/part_csv')

# save csv and zip
pred_scores, _, _, _ = predict(model, data_set['val'], data_loader['val'])
val_pred = gen_submission(val_pd[['ImageName', 'AttrKey']], pred_scores, catidx_map=data_set['val'].catidx_map)

if len(select_AttrIdx) < 8:
    val_pred[['ImageName', 'AttrKey', 'AttrValueProbs']].to_csv('./val_pred2/part_csv/%s.csv' % model_name,
                                                                 header=None, index=False)
else:
    val_pred[['ImageName', 'AttrKey', 'AttrValueProbs']].to_csv('./val_pred2/csv/%s.csv' % model_name,
                                                                 header=None, index=False)

print val_pred[['ImageName', 'AttrKey', 'AttrValueProbs']].info()
